#!/usr/bin/env python
# Python script to export data from simpledb
import os
import sys
import threading
import boto
from boto.sdb.item import Item
from boto.sdb.connection import SDBConnection
from time import strftime

class ItemThread(threading.Thread):
	'''
		The ItemThread class is a threaded class that fetches multiple items from SimpleDB
		at the same time. Each Thread connects to the sdb domain and fetches the next 'x'
		quantity of items from the domain. At the end, each thread is bound to the calling
		thread and unified into a single list of Item objects.
	'''
	def __init__(self, name, domain_name, item_names):
		threading.Thread.__init__(self, name=name)
		self.domain_name = domain_name
		self.conn = SDBConnection()
		self.item_names = item_names
		self.items = []

	def run(self):
		'''
			This method loops through a list of Item names and populates the shell item
			with valid content. Once the Item is populated, it gets placed in the list to
			be passed back into the calling thread.
		'''
		for item_name in self.item_names:
			item = self.conn.get_attributes(self.domain_name, item_name)
			self.items.append(item)

class sdbExport(SDBConnection):
	'''
		The sdbExport class simply extends the SDBConnection class defined in
		boto.sdb.connection and adds a method to run threaded queries properly.
		All other methods are leaved untouched to allow us to use it for all
		other standard queries.
	'''
	def threaded_query(self, domain_or_name, query='', max_items=None, next_token=None, num_threads=6):
		'''
			Returns a list of fully populated items that match the query provided.

			The name/value pairs for all of the matching item names are retrieved in a number of separate
			threads (specified by num_threads) to achieve maximum throughput.
			The ResultSet that is returned has an attribute called next_token that can be used
			to retrieve additional results for the same query.
		'''
		# Set a hard limit on the quantity of threads we are allowed using
		if num_threads >= 81:
			num_threads = 80
			# If we have the verbose flag set, display a message
			if options.is_verbose:
				print "Threads originally set to %s. We only allow a maximum of 80"

		domain, domain_name = self.get_domain_and_name(domain_or_name)
		if max_items and num_threads > max_items:
			num_threads = max_items
		rs = self.query(domain_or_name, query, max_items, next_token)
		threads = []
		n = len(rs) / num_threads
		for i in range(0, num_threads):
			if i+1 == num_threads:
				thread = ItemThread('Thread-%d' % i, domain_name, rs[n*i:])
			else:
				thread = ItemThread('Thread-%d' % i, domain_name, rs[n*i:n*(i+1)])
			threads.append(thread)
			thread.start()
		del rs[0:]
		for thread in threads:
			thread.join()
			for item in thread.items:
				rs.append(item)
		return rs

def getDomains():
	'''
		The get_domains function fetches a list of all the domains from the simpledb connection.
		Once is has the list, it checks if the user used the -d flag. If so, it queries the given
		domain to fetch the list of returned results. If not, it simply returns the list of domains
		to the user and exits successfully.
	'''
	# If we only want to list the domains, do so, then exit
	if options.list_domains:
		try:
			# If we have verbose reporting, place verbose msg
			if options.is_verbose:
				print "Domain list:"

			# Loop through the domain lists we get and display them
			for d in sdb.get_all_domains():
				print str(d)[7: ]
			sys.exit(0)
		except boto.exception.SDBResponseError:
			print "Error fetching domain list. Please check your authentication options."
			# If we have the verbose flag set, display the exception message
			if options.is_verbose:
				print sys.exc_info()[0]
			sys.exit(1)
	else:
		try:
			d = sdb.get_domain(options.domain, validate=True)
			getItems(d)
		except boto.exception.SDBResponseError:
			print "Domain doesnt exist. Please use the -l option to list valid domains"
			# If we have the verbose flag set, display the exception message
			if options.is_verbose:
				print sys.exc_info()[0]
			sys.exit(1)

def getItems(d):
	'''
		The get_items function is called by the get_domains function with a given domain object.
		If the user specified the -f option, it opens/creates a file specified by the -f option
		and starts iterating through the domain items by creating -t threads and inserts them
		into the file. If the user didnt set the -f flag, it simply returns the content to the
		stdout.
	'''
	# If the user used the -f or --file option, then save the data into a file, else print to stdout
	if options.filename:
		# Create a file to write the data to
		f = open(options.filename,'w')
		for i in parseItem(d):
			data = i
			data['ItemName'] = i.name
			f.write(str(data)+"\n")
		f.flush()
		f.close()
	else:
		try:
			for i in parseItem(d):
				data = i
				data['ItemName'] = i.name
				print data
			sys.exit(0)
		except boto.exception.SDBResponseError:
			print "Error processing domian."
			# If we have the verbose flag set, display the exception message
			if options.is_verbose:
				print sys.exc_info()[1]
			sys.exit(1)

def parseItem(domain, query=''):
	'''
		The parse_item function loops through a token list of items within the domain returning
		the list of item data to the get_items function.
	'''
	global next_token
	more_results = True
	num_results = 0
	next_token = None

	# Set the starting token if we need to. We need to remove the starting token once set, or we will loop
	if options.token:
		f = open(options.token,'r')
		next_token = f.read()
		f.close()
		options.token = None

	while more_results:
		rs = domain.connection.threaded_query(domain.name, query, None, next_token, options.threads)
		for item in rs:
			if options.max:
				if num_results == options.max:
					raise StopIteration
			yield item
			num_results += 1
		next_token = rs.next_token
		more_results = next_token != None

if __name__ == "__main__":
	# Parse our command line options
	from optparse import OptionParser
	parser = OptionParser(version="%prog R12")

	# Define our options
	parser.usage = "%prog [options]"
	parser.add_option("-i", "--id", type="string", dest="username", help="set the AWS_ACCESS_KEY_ID variable", metavar="ID")
	parser.add_option("-k", "--key", type="string", dest="password", help="set the AWS_SECRET_ACCESS_KEY variable", metavar="KEY")
	parser.add_option("-v", "--verbose", action="store_true", dest="is_verbose", default=False, help="print additional information")
	parser.add_option("-l", "--list", action="store_true", dest="list_domains", default=False, help="fetch a list of the domains")
	parser.add_option("-d", "--domain", type="string", dest="domain", help="fetch data from a domain")
	parser.add_option("-f", "--file", type="string", dest="filename", help="output the data into the specified file")
	parser.add_option("-t", "--threads", type="int", dest="threads", default=10, help="set the quantity of threads to run")
	parser.add_option("-n", "--nexttoken", type="string", dest="token", default=None, help="set the start token file", metavar="TOKEN")
	parser.add_option("-m", "--maxnumberofitems", type="int", dest="max", default=None, help="set the maximum rows to return")
	(options, args) = parser.parse_args()

	# Try set the environment variables used for authentication
	if options.username:
		os.environ['AWS_ACCESS_KEY_ID'] = options.username
	if options.password:
		os.environ['AWS_SECRET_ACCESS_KEY'] = options.password

	# Check to make sure we have valid usernames and passwords
	try:
		sdb = sdbExport()
	except:
		print "Error connecting to SimpleDB. Please check your access key and password. Use --help for help."
		sys.exit(1)

	# Set the default action to list domains
	if not options.domain:
		options.list_domains=True

	# Get the domain list
	try:
		next_token = None
		getDomains()
	except (KeyboardInterrupt, SystemExit):
		if next_token:
			# Create a file to write the last token into
			f = open("token"+strftime("_%Y%m%d%H%M%S"),'w')
			f.write(next_token)
			f.flush()
			f.close()

